import sys
import os
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QPushButton, QTextEdit
from PyQt5.QtGui import QPixmap, QImage, QColor
from PyQt5.QtCore import Qt
from PIL import Image


class ImageViewer(QWidget):
    def __init__(self, parent_folder):
        super().__init__()
        self.parent_folder = parent_folder
        self.image_names = {}
        self.current_index = 0
        self.drag_start_pos = None

        self.init_ui()

    def init_ui(self):
        self.setWindowTitle('Image Viewer')
        self.setFixedSize(1000, 800)  # 设置固定大小
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)

        self.image_label = QLabel()
        self.layout.addWidget(self.image_label)

        self.info_label = QTextEdit()
        self.info_label.setReadOnly(True)
        self.layout.addWidget(self.info_label)

        self.image_name_label = QLabel()  # 显示图片名称的标签
        self.layout.addWidget(self.image_name_label)

        self.prev_button = QPushButton('上一张')
        self.prev_button.clicked.connect(self.show_prev_image)
        self.layout.addWidget(self.prev_button)

        self.next_button = QPushButton('下一张')
        self.next_button.clicked.connect(self.show_next_image)
        self.layout.addWidget(self.next_button)

        self.load_images()
        self.show_image()

        # 启用鼠标跟踪
        self.setMouseTracking(True)

    def load_images(self):
        for folder_name in os.listdir(self.parent_folder):
            folder_path = os.path.join(self.parent_folder, folder_name)
            if os.path.isdir(folder_path):
                folder_path=os.path.join(folder_path,"predictResult")
                if not os.path.exists(folder_path):
                    continue
                for file_name in os.listdir(folder_path):
                    if file_name.endswith('.jpg') or file_name.endswith('.png'):
                        file_path = os.path.join(folder_path, file_name)
                        self.image_names.setdefault(file_name, []).append(file_path)

        # 移除只存在于一个文件夹中的图片
        self.image_names = {file_name: file_paths for file_name, file_paths in self.image_names.items() if
                            len(file_paths) > 1}

    def show_image(self):
        if not self.image_names:
            return

        file_name = list(self.image_names.keys())[self.current_index]
        file_paths = self.image_names[file_name]
        file_path = file_paths[0]
        print(file_path)

        # 加载图像
        image = QImage(file_path)
        if image.isNull():
            print("Image is null")
            return  # 如果图像加载失败，直接返回

        # 检查图像的通道数
        if image.isGrayscale():
            print("Image is single-channel (monochrome)")
            # 如果图像是单通道的，转换为RGB格式
            image = image.convertToFormat(QImage.Format_RGB888)

            # 为了正确处理单通道图像，将其转换为伪彩色图像
            color_table = [QColor(i, i, i).rgb() for i in range(256)]
            image.setColorTable(color_table)

        # 将图像转换为pixmap
        pixmap = QPixmap.fromImage(image)

        # 设置图像到 label 中
        self.image_label.setPixmap(pixmap)
        self.image_label.setScaledContents(True)  # 图片自适应窗口大小

        # 将所有路径显示在 info_label 中
        self.info_label.setText('\n'.join(file_paths))

        # 显示图片名称并添加颜色
        colored_text = f"<font color='blue'>{file_name}</font>"
        self.image_name_label.setText(colored_text)

    def show_prev_image(self):
        if self.current_index > 0:
            self.current_index -= 1
            self.show_image()

    def show_next_image(self):
        if self.current_index < len(self.image_names) - 1:
            self.current_index += 1
            self.show_image()

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.drag_start_pos = event.globalPos()

    def mouseMoveEvent(self, event):
        if event.buttons() == Qt.LeftButton:
            if self.drag_start_pos:
                new_pos = self.mapToGlobal(event.pos() - self.drag_start_pos)
                self.move(new_pos)

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.drag_start_pos = None

    def resizeEvent(self, event):
        super().resizeEvent(event)
        self.show_image()


if __name__ == '__main__':
    app = QApplication(sys.argv)

    # 父文件夹路径
    parent_folder = '/home/jolly/develop/PycharmProjects/remote_sensor/bijie256'
    viewer = ImageViewer(parent_folder)
    viewer.show()

    sys.exit(app.exec_())